from paddleocr import PaddleOCR
import os
from colorama import Fore
import time


def ocr_recognition(input_dir, output_dir, callback):
    """
    使用 PaddleOCR 对指定目录中的图像文件进行文字识别，并将识别结果保存到文本文件中。

    参数:
    - input_dir (str): 输入文件目录路径。
    - output_dir (str): 输出文件目录路径，用于保存识别结果文本文件。
    - callback (function): 回调函数，用于处理识别结果。

    返回:
    无，直接调用回调函数处理识别结果。
    """
    # 记录程序开始时间
    start_time = time.time()

    # 初始化 PaddleOCR 模型
    ocr = PaddleOCR(use_angle_cls=True, lang="ch")

    # 确保输出目录存在，如果不存在则创建
    os.makedirs(output_dir, exist_ok=True)

    # 统计文件的数量
    total_files = 0
    processed_files = 0
    skipped_files = 0

    # 遍历文件目录及其子目录
    for root, dirs, files in os.walk(input_dir):
        # 获取相对于输入目录的路径
        relative_path = os.path.relpath(root, input_dir)
        # 构造对应的输出目录路径
        output_subdir = os.path.join(output_dir, relative_path)
        # 确保输出子目录存在，如果不存在则创建
        os.makedirs(output_subdir, exist_ok=True)

        # 遍历当前目录下的所有文件
        for filename in files:
            # 拼接输入文件的完整路径
            input_file = os.path.join(root, filename)
            # 更新文件总数
            total_files += 1

            # 判断文件扩展名是否为支持 OCR 的格式
            supported_formats = ['.png', '.jpg', '.jpeg', '.pdf']
            ext = os.path.splitext(filename)[1]
            if ext.lower() in supported_formats:
                # 构造输出文件名 文件名.后缀名.txt
                output_filename = filename + ".txt"
                output_file = os.path.join(output_subdir, output_filename)

                # 检查是否已存在相应的文本文件，如果存在说明已经识别过，跳过当前文件
                if os.path.exists(output_file):
                    skipped_files += 1
                    print(Fore.RED + f"==> 跳过文件 {input_file}: 文件已经识别")
                    continue

                # 打印绿色提示
                print(Fore.BLUE + f"==> 开始识别文件 {input_file}")

                # 使用 PaddleOCR 模型进行文字识别
                result = ocr.ocr(input_file, cls=True)

                # 将识别结果拼接成一个字符串
                text_result = ""
                for idx in range(len(result)):
                    res = result[idx]
                    for line in res:
                        text_result += line[1][0] + "\n"

                # 将识别结果写入到文本文件中
                with open(output_file, 'w', encoding='utf-8') as f:
                    f.write(text_result)

                print(Fore.GREEN + f"==> 文件 {filename} 识别完成，识别结果已写入 {output_file}")
                # 更新已处理的文件数量
                processed_files += 1
                # 调用回调函数并传递识别结果
                formatted_word_boxes = []
                for res in result:
                    for line in res:
                        bbox = line[0]
                        formatted_bbox = [
                            min(bbox[0][0], bbox[3][0]),  # left
                            min(bbox[0][1], bbox[1][1]),  # top
                            max(bbox[1][0], bbox[2][0]),  # right
                            max(bbox[2][1], bbox[3][1])   # bottom
                        ]
                        formatted_word_boxes.append(
                            [line[1][0], formatted_bbox])
                callback(input_file, formatted_word_boxes)
            else:
                # 更新跳过的文件数量
                skipped_files += 1
                print(Fore.RED + f"==> 跳过文件 {filename}: 文件类型不支持 OCR")

    # 计算程序执行时间
    end_time = time.time()
    total_time_seconds = end_time - start_time

    # 输出程序执行情况及总耗时
    print(Fore.YELLOW +f"==> 文件总数： {total_files} ，识别数：{processed_files} ，跳过：{skipped_files} ")
    print(Fore.YELLOW + f"==> 本次运行总耗时：{total_time_seconds:.2f} 秒")
